package soundsourceloc;

import gnu.io.CommPortIdentifier;
import gnu.io.PortInUseException;
import gnu.io.SerialPort;
import gnu.io.UnsupportedCommOperationException;

import java.awt.BasicStroke;
import java.awt.BorderLayout;
import java.awt.Dimension;
import java.awt.Graphics;
import java.awt.Graphics2D;
import java.awt.event.WindowAdapter;
import java.awt.event.WindowEvent;
import java.io.IOException;
import java.io.OutputStream;
import java.util.Enumeration;
import java.util.logging.Level;
import java.util.logging.Logger;

import javax.sound.sampled.AudioFormat;
import javax.sound.sampled.AudioInputStream;
import javax.sound.sampled.AudioSystem;
import javax.sound.sampled.TargetDataLine;
import javax.swing.JFrame;
import javax.swing.JPanel;

import org.tritonus.share.sampled.FloatInputStream;
import org.tritonus.share.sampled.FloatSampleBuffer;

/**
 * This class computes the direction of the source of the sound it hears.
 * 
 * It uses 2 microphones, and compute the time of arrival difference of sound
 * between them to estimate the sound source localization.
 * 
 * The direction of the sound is then drawn on screen, and used to control a
 * robot head pan servo (through a serial command sent to an Arduino board).
 * 
 * Packaging note: this class uses some utilities class from tritonus.org. The
 * corresponding sources are included in this distrib, as no pre-compiled
 * version of this tritonus version exist (last CVS content), and compiling
 * Tritonus in general is quite messy...
 * 
 * @author Frederic Pesquet (fpesquet at gmail dot com)
 */
public class SoundSourceLoc extends JFrame {
	private static final long serialVersionUID = 1L;

	protected static final Logger LOG = Logger.getLogger(SoundSourceLoc.class
			.getName());

	/**
	 * Name of the serial port used to talk to the Arduino controller which
	 * handles the robot head pan servo
	 */
	static private final String _arduinoSerialPortName = "/dev/ttyUSB0";

	/**
	 * Max time shift between right and left mic in number of samples. <br/>
	 * This typically depends on the sample rate and the distance between
	 * microphones. <br/>
	 * You can either compute this with clever formulas involving sound speed
	 * and microphones distance, or just try and put the max value you get with
	 * extreme loc of sound. Guess what I did :-)
	 */
	static private final int _nbSamplesMaxDiff = 15;

	/**
	 * Buffer size on which we will try to locate sound. <br/>
	 * This is a number of samples, and depends on sample rate, and speed of
	 * sound loc change we want to detect. Lower values mean compute sound loc
	 * often, but accuracy is quite low as we compute on a very small slice of
	 * sound. <br/>
	 * Empirically, I found that computing on long sounds is better, here 4096
	 * samples at 44 KHz sampling rate means about one second of sound => we
	 * reevaluate sound loc every second. <br/>
	 * Notice that the larger the value, the most computation we do, as we time
	 * shift on the whole buffer.
	 */
	static private final int _bufferSize = 4096;

	/**
	 * sound speed in meters per seconds
	 */
	static private final float _soundSpeed = 344;

	/**
	 * sound sampling rate in Hz
	 */
	static private final int _soundSamplingRate = 44100;

	/**
	 * min time to move again after first move, in ms
	 * */
	static private final long _minTimeBetweenMoves = 1000;

	/**
	 * Distance between microphones in meters
	 */
	static private final float _distanceBetweenMicrophones = 0.1f;

	/**
	 * Take a point for sound loc is level > 105% of mean level. <br/>
	 * This allows to compute sound loc only for "meaningful" sounds, not
	 * background noise.
	 */
	static private float _minLevelFactorForValidLoc = 1.05f;

	/** An utility to compute the running average of sound power */
	private final RunningAverage _averageSoundLevel = new RunningAverage(50);

	/** the panel that draw the last sound localization as an arc */
	private final SoundLocDraw _soundLocDraw;

	/** sound samples input stream */
	private final FloatInputStream _inputStream;

	/** sound samples input buffer */
	private final FloatSampleBuffer _buffer;

	/** serial port to arduino board */
	private SerialPort _arduinoSerialPort;

	/** serial stream */
	private OutputStream _arduinoSerialPortOutput;
	// private InputStream _arduinoSerialPortInput;

	/** last time we moved (absolute time in system ms) */
	private long _lastTimeMoveDone = System.currentTimeMillis();

	public SoundSourceLoc() throws Exception {
		super("Sound Source Localization");

		_soundLocDraw = new SoundLocDraw();
		getContentPane().add(_soundLocDraw, BorderLayout.CENTER);

		// sample 2 channels, 16 bits.
		AudioFormat samplingFormat = new AudioFormat(_soundSamplingRate, 16, 2,
				true, false);
		TargetDataLine inputMic = AudioSystem.getTargetDataLine(samplingFormat);
		inputMic.open();
		inputMic.start();
		if (!inputMic.isActive()) {
			throw new Exception("input mic not ready, audio issue?");
		}
		AudioInputStream rawIn = new AudioInputStream(inputMic);
		_inputStream = new FloatInputStream(rawIn);
		_buffer = new FloatSampleBuffer(2, _bufferSize, _inputStream
				.getSampleRate());

		connectToArduinoCtrl();

		addWindowListener(new WindowAdapter() {
			@Override
			public void windowClosed(WindowEvent e) {
				shutdown();
			}
		});
	}

	private void shutdown() {
		disconnectFromArduinoCtrl();
	}

	@SuppressWarnings("unchecked")
	private void connectToArduinoCtrl() throws IOException, PortInUseException,
			UnsupportedCommOperationException {
		Enumeration<CommPortIdentifier> portList = CommPortIdentifier
				.getPortIdentifiers();
		while (portList.hasMoreElements()) {
			CommPortIdentifier portId = portList.nextElement();
			if (portId.getName().equals(_arduinoSerialPortName)) {
				LOG.info("open_port:" + portId.getName());
				_arduinoSerialPort = (SerialPort) portId.open("Arduino Serial",
						2000);
				// _serialPort.setFlowControlMode(SerialPort.FLOWCONTROL_NONE);
				// _arduinoSerialPortInput =
				// _arduinoSerialPort.getInputStream();
				_arduinoSerialPortOutput = _arduinoSerialPort.getOutputStream();
				_arduinoSerialPort.setSerialPortParams(9600, 8,
						SerialPort.STOPBITS_1, SerialPort.PARITY_NONE);
			}
		}
		if (_arduinoSerialPort == null) {
			throw new IOException("Can't connect to Arduino");
		}
	}

	private void disconnectFromArduinoCtrl() {
		try {
			if (_arduinoSerialPort != null) {
				LOG.log(Level.INFO, "closing serial port");
				_arduinoSerialPort.close();
			}
		} catch (Exception e) {
			LOG.log(Level.WARNING, "can't close serial port", e);
		}
		_arduinoSerialPort = null;
	}

	/**
	 * Set desired head pan position
	 * 
	 * @param anglerad
	 *            desired angle, in radians
	 */
	private void setHeadPanTo(double anglerad) {
		try {
			int servoAngle = (int) Math.toDegrees(anglerad) + 90;
			if (servoAngle < 0)
				servoAngle = 0;
			if (servoAngle > 180)
				servoAngle = 180;
			LOG.log(Level.INFO, "set angle pan head to " + servoAngle);
			_arduinoSerialPortOutput.write('S');
			_arduinoSerialPortOutput.write('P');
			_arduinoSerialPortOutput.write(servoAngle);
			_arduinoSerialPortOutput.flush();
		} catch (Exception e) {
			LOG.log(Level.WARNING, "Error occurred writing to arduino", e);
		}
	}

	/**
	 * Main loop: read a buffer, compute sound source localization, iterate.
	 */
	public void run() {
		while (!_inputStream.isDone()) {
			processNextSoundBlock();
		}
	}

	/**
	 * This is the core of the sound source localization: it takes the
	 * right/left sampled sounds, and compute their differences while delaying
	 * one channel more and more.<br/>
	 * => the delay for which the difference is minimal is the real delay
	 * between the right/left sounds, from which we can deduce the sound source
	 * localization
	 */
	private void processNextSoundBlock() {
		_inputStream.read(_buffer);
		float minDiff = Float.MAX_VALUE;
		int minDiffTime = -1;
		float[] right = _buffer.getChannel(0);
		float[] left = _buffer.getChannel(1);

		// 'slide' time to find minimum of right/left sound differences
		for (int t = -_nbSamplesMaxDiff; t < _nbSamplesMaxDiff; t++) {
			// compute sum of differences as the cross-correlation-like measure:
			float diff = 0;
			for (int i = _nbSamplesMaxDiff; i < left.length - _nbSamplesMaxDiff; i++) {
				diff += Math.abs(left[i] - right[i + t]);
			}
			if (diff < minDiff) {
				minDiff = diff;
				minDiffTime = t;
			}
		}
		// LOG.info("mean diff time is " + minDiffTime);

		// compute the sound level (i.e. "loudness" of the sound):
		float level = computeLevel(right, left);
		// update the average sound level with this new measure:
		_averageSoundLevel.newValue(level);
		// relative sound level of this sample compared to average:
		float relativeLevel = level / _averageSoundLevel.getMean();
		// float relativeLevel = level / 0.088f;

		// if sound is lound enough, and not an extreme (=usually false
		// measure), then draw it:
		if ((relativeLevel > _minLevelFactorForValidLoc)
				&& (minDiffTime > -_nbSamplesMaxDiff)
				&& (minDiffTime < _nbSamplesMaxDiff)) {
			// angle of sound loc is approx. linear with time shift, and between
			// -90 deg - + 90 deg:
			// float angle = -(float) (Math.PI / 2 * ((float) minDiffTime /
			// (float) _nbSamplesMaxDiff));

			// real computation of angle depending on diff time, sampling rates,
			// and geometry
			// (thanks Mathieu from Pobot :-) ):
			float angle = -(float) Math.asin((minDiffTime * _soundSpeed)
					/ (_soundSamplingRate * _distanceBetweenMicrophones));

			// as we move the servo head and it makes sound, this
			// disturb the localization process and results in oscillations
			// and bad detections after the first one.
			// => to avoid this, stop reacting for some time once we have
			// decided for a direction
			if (System.currentTimeMillis() > _lastTimeMoveDone
					+ _minTimeBetweenMoves) {
				_lastTimeMoveDone = System.currentTimeMillis();
				_soundLocDraw.setSound(angle, relativeLevel);
				setHeadPanTo(angle);
			}
		}
	}

	/**
	 * Compute avererage sound level (i.e. power) for left/right channels.
	 * 
	 * Notice we could probably do the computation on some samples only (for
	 * example one over 4 samples) without loosing much accuracy here. This
	 * would reduce computation time.
	 */
	private float computeLevel(float[] right, float[] left) {
		float level = 0;
		for (int i = 0; i < right.length; i++) {
			float s = (left[i] + right[i]) / 2;
			level += (s * s);
		}
		level /= right.length;
		level = (float) Math.sqrt(level);
		return level;
	}

	/**
	 * An utility class to compute the running average of a signal
	 */
	static public class RunningAverage {
		private final int _nbValuesForAverage;
		private int _nbValues;
		private float _mean;

		RunningAverage(int nbValuesForAverage) {
			_nbValuesForAverage = nbValuesForAverage;
		}

		public void newValue(float v) {
			if (_nbValues < _nbValuesForAverage) {
				_nbValues++;
			} else {
				LOG.fine("average sound level is " + _mean);
			}
			_mean = ((_mean * (_nbValues - 1)) + v) / _nbValues;
		}

		public float getMean() {
			return _mean;
		}
	}

	/**
	 * A simple panel that draws the sound source localization angle, with a
	 * thickness dependending on the sound level.
	 */
	@SuppressWarnings("serial")
	static private class SoundLocDraw extends JPanel {
		// sound angle, between -PI/2...+PI/2
		private float _angle;

		// relative power with respect to mean power (1.0=mean power)
		private float _relativePower;

		public void setSound(float angle, float relativePower) {
			_angle = angle;
			_relativePower = relativePower;
			repaint();
		}

		@Override
		protected void paintComponent(Graphics g) {
			super.paintComponent(g);
			Graphics2D g2d = (Graphics2D) g;
			Dimension d = getSize();
			int radius = Math.min(d.height, d.width / 2);
			int cx = d.width / 2;
			int cy = 0;
			int tx = cx + (int) (Math.cos(_angle + Math.PI / 2) * radius);
			int ty = cy + (int) (Math.sin(_angle + Math.PI / 2) * radius);
			g2d.drawOval(cx - radius, cy - radius, radius * 2, radius * 2);
			// use larger strokes for louder sounds:
			g2d.setStroke(new BasicStroke(1 + (int) ((Math.max(_relativePower,
					1) - 1.0) * 10)));
			g2d.drawLine(cx, cy, tx, ty);
		}
	}

	/**
	 * Entry point: create the frame, and start listening to sound until closed.
	 */
	public static void main(String[] args) throws Exception {
		SoundSourceLoc snd = new SoundSourceLoc();
		snd.setSize(800, 400);
		snd.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
		snd.setLocationRelativeTo(null);
		snd.setVisible(true);
		snd.run();
	}
}
